package cc.spider.web.queue;

import cc.spider.core.queue.AbsUniqueQueue;
import cc.spider.core.queue.DefaultTaskQueue;
import cc.spider.core.queue.PriorityQueue;
import cc.spider.model.TaskInfo;
import cc.spider.web.model.http.Request;
import org.apache.commons.lang3.StringUtils;

import java.util.HashSet;
import java.util.Set;

/**
 * 具有优先级特性的去重队列，通过URL判断是否重复。默认POST模式下不去去重。
 */
public class UniquePriorityQueue extends AbsUniqueQueue {
    PriorityQueue taskQueue = new PriorityQueue();
    Set<String> urls = new HashSet<>();

    @Override
    public boolean isUnique(TaskInfo taskInfo) {
        return (taskInfo instanceof Request && StringUtils.containsIgnoreCase("POST", ((Request) taskInfo).getMethod())) || urls.add(taskInfo.getUuKey());
    }

    @Override
    public boolean enqueueWhenUnique(TaskInfo taskInfo) {
        if(taskQueue.enqueue(taskInfo)){
            return true;
        }else{
            urls.remove(taskInfo.getUuKey());
            return false;
        }
    }

    @Override
    public TaskInfo dequeue() {
        return taskQueue.dequeue();
    }
}
